from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from pydantic import BaseModel, Field
from typing import List
from langchain_core.tools import tool
from typing import Annotated
from utils import websearch
from langgraph.prebuilt import create_react_agent


@tool
def web_search(query: Annotated[str, "互联网查询内容"]):
    """
    通过web_search工具查询互联网上的信息
    """
    return websearch(query)


_executor_system_template = """
您是一个优秀的子任务执行者，你可以根据子任务的名称及历史参考信息完成子任务的信息查询
你可以使用一下工具帮助你更好的完成该任务:
{tools_name}
"""
_executor_human_template = """
参考信息:
{infos}
子任务名称:
{task}

"""


class Executor:
    def __init__(self, llm):

        _tools = [
            web_search,
        ]
        _prompt = ChatPromptTemplate.from_messages(
            [
                ("system", _executor_system_template),
                ("human", _executor_human_template),
            ]
        )
        _prompt = _prompt.partial(tools_name=",".join([_tool.name for _tool in _tools]))

        _llm_with_tools_agent = create_react_agent(llm, tools=_tools)

        self._chain = _prompt | _llm_with_tools_agent
        self._parser = StrOutputParser()

    def __call__(self, state):
        _rt = self._chain.invoke(state)
        _messages = _rt["messages"]

        return self._parser.invoke(_messages[-1])


if __name__ == "__main__":
    from langchain_openai import ChatOpenAI

    _llm = ChatOpenAI(
        base_url="http://192.168.10.11:60026/v1",
        model="qwen2.5:7b",
        api_key="ollama",
        temperature=0.4,
    )
    executor = Executor(_llm)
    _rt = executor({"infos": [], "task": "确定巴黎奥运会女子十米跳台冠军"})
    print(_rt)
